from __future__ import annotations

from typing import (
    Optional,
    Union,
    Dict,
    NamedTuple,
    Tuple,
    List,
)
from typing_extensions import Literal
import numpy.typing as npt
from common.typing_utils import numeric

import json
from math import floor, ceil

import affine
from pyproj import CRS

from common.logger import logger
from common.constants.coord_sys import NOT_COMPATIBLE_SRS_MSG, COMPATIBLE_EPSG
from common.geometry.bbox import BoundingBox
from common.geometry.geom import Geom, bbox_to_geom
from common.geometry.tools import project_geom
from common.geometry.srs import wkt_to_epsg


class ValueInterpretationMeta(NamedTuple):
    """Class for metadata representing how to interpret value in a band.

    Parameters
    ----------
    scale: Optional[Union[numeric, npt.DTypeLike]]
        Scale factor for that band.
    offset: Optional[Union[numeric, npt.DTypeLike]]
        Offset value for that band.
    nodata: Optional[Union[numeric, npt.DTypeLike]]
        Nodata value for that band.
    """

    scale: Optional[Union[numeric, npt.DTypeLike]]
    offset: Optional[Union[numeric, npt.DTypeLike]]
    nodata: Optional[Union[numeric, npt.DTypeLike]]


class SceneMeta(NamedTuple):
    """Class for metadata describing a Scene.

    Parameters
    ----------
    n_cols: int
        Number of columns in the scene.
    n_rows: int
        Number of rows in the scene.
    n_bands: int
        Number of bands in the scene.
    prj_wkt: str
        Projection as a WKT string for that scene.
    geotrans: Tuple[float, float, float, float, float, float]
        Geotransformation as a 6-element tuple for that scene.
    value_interpretations: Optional[List[ValueInterpretationMeta]]
        List of metadata for interpreting values in each band, default is None
    """

    n_cols: int
    n_rows: int
    n_bands: int
    prj_wkt: str
    geotrans: Tuple[float, float, float, float, float, float]
    value_interpretations: Optional[List[ValueInterpretationMeta]]

    def is_same_spatial_domain(self, other: SceneMeta):
        flag = (
            self.n_rows == other.n_rows
            and self.n_cols == other.n_cols
            and self.prj_wkt == other.prj_wkt
            and tuple(self.geotrans) == tuple(other.geotrans)
        )
        return flag

    @property
    def n_pixels(self) -> int:
        """number of pixels in the scene, as n_rows X n_cols"""
        return self.n_rows * self.n_cols

    @property
    def n_size(self) -> int:
        """size of the scene, as n_rows X n_cols X n_bands"""
        return self.n_rows * self.n_cols * self.n_bands

    @property
    def x_res(self) -> numeric:
        """resolution in the west-to-east direction"""
        return self.geotrans[1]

    @property
    def y_res(self) -> numeric:
        """resolution in the north-to-south direction"""
        return self.geotrans[5]

    @property
    def bounds(self) -> BoundingBox:
        """bounding box for the scene represented by the `SceneMeta`"""
        return calc_bounds_from_geotrans(self.geotrans, self.n_rows, self.n_cols)

    def geom(self, enforce_wgs84=True) -> Geom:
        """bounding box of the scene as a geometry

        Parameters
        ----------
        enforce_wgs84 : bool, optional
            Whether to enforce the returned geometry to be under WGS84 coordinate system,
            by default True to enforce,
            otherwise to keep in the scene's own coordinate system.

        Returns
        -------
        Geom
            The result geometry.
        """
        geom = Geom.from_bbox(
            calc_bounds_from_geotrans(self.geotrans, self.n_rows, self.n_cols)
        )
        if enforce_wgs84:
            geom = project_geom(geom, src_epsg=self.prj_epsg, dest_epsg="EPSG:4326")
        return geom

    def intersects(self, another_scene_meta: SceneMeta) -> bool:
        """Whether an area represented by a SceneMeta is intersected with another one"""
        if self.prj_epsg_no == another_scene_meta.prj_epsg_no:
            geom0 = self.geom(enforce_wgs84=False)
            geom1 = another_scene_meta.geom(enforce_wgs84=False)
        else:
            geom0 = self.geom(enforce_wgs84=True)
            geom1 = another_scene_meta.geom(enforce_wgs84=True)
        return geom0.intersects(geom1)

    @property
    def prj_epsg(self) -> str:
        """projection as an EPSG string"""
        return wkt_to_epsg(self.prj_wkt, get_str=True)

    @property
    def prj_epsg_no(self) -> str:
        """projection as an EPSG number"""
        return wkt_to_epsg(self.prj_wkt, get_str=False)

    @property
    def affine(self) -> affine.Affine:
        """geotransformation as an affine matrix"""
        return affine.Affine.from_gdal(*self.geotrans)

    def to_dict(self) -> Dict:
        """Export the `SceneMeta` object to a dictionary."""
        return self._asdict()

    @staticmethod
    def from_dict(input_dict: Dict) -> SceneMeta:
        """Create a `SceneMeta` object from a given input dictionary."""
        if (
            "value_interpretations" in input_dict
            and input_dict["value_interpretations"] is not None
        ):
            vis = []
            for vi in input_dict["value_interpretations"]:
                vis.append(ValueInterpretationMeta(*vi))
            input_dict["value_interpretations"] = vis
        return SceneMeta(**input_dict)

    def to_json(self) -> str:
        """Export the `SceneMeta` object to a json string."""
        return json.dumps(self._asdict())

    @staticmethod
    def from_json(json_str: str) -> SceneMeta:
        """Create a `SceneMeta` object from a given json string."""
        info_meta = json.loads(json_str)
        return SceneMeta.from_dict(info_meta)

    def update(self, **kwargs) -> SceneMeta:
        """Update and create a new SceneMeta object by replacing fields in the original SceneMeta
        using information provided in kwargs.

        Parameters
        ----------
        n_cols: int
            Number of columns in the scene.
        n_rows: int
            Number of rows in the scene.
        n_bands: int
            Number of bands in the scene.
        prj_wkt: str
            Projection as a WKT string for that scene.
        geotrans: Tuple[float, float, float, float, float, float]
            Geotransformation as a 6-element tuple for that scene.
        value_interpretations: Optional[List[ValueInterpretationMeta]]
            List of metadata for interpreting values in each band, default is None
        """
        dict_meta = self._asdict()
        for key, value in kwargs.items():
            if key not in dict_meta:
                logger.warning(f"key {key} is not a supported one, skipped.")
            else:
                dict_meta[key] = value
        res = SceneMeta(**dict_meta)
        return res

    def copy(self) -> SceneMeta:
        """Make a copy of the `SceneMeta` object."""
        return self.update()

    def scale_update(
        self,
        level: numeric,
        boundary_treatment: Literal["round", "shrink", "expand"] = "round",
    ) -> SceneMeta:
        """Update and create a new `SceneMeta` object with up or down scale but still representing the same area

        Parameters
        ----------
        level: numeric
            The level used to change the scale.
            See `Notes` for more details.

        boundary_treatment: Literal["round", "shrink", "expand"]
            How to treat boundary condition, should be one of "round", "shrink", or "expand".
            By default "round". See `Notes` for more details.

        Returns
        -------
        SceneMeta
            A new `SceneMeta` object with updated scale

        Notes
        -----
        If `level` is 1, then no actual change; otherwise, the resulting `SceneMeta` object would represent
        the same area but with changed resolution as `original_resolution / level`.
        For example, if the original resolution is 10m, when `level = 2` is given,
        the resulting `SceneMeta` would have a resolution of 5m;
        and when the `level = 0.5` is given, the resulting `SceneMeta` would have a resolution of 20m.

        The method try to keep the resulting meta having the same upper and left coordinates.
        For the bottom and right coordinates, the behavior is controlled by the `boundary_treatment` argument.
        If "round", the method would round the number of columns or number of rows computed.
        If "shrink", it would shrink the number of columns or rows to only include fully valid zones.
        If "expand, it would expand the number of columns or rows to include zones that could have valid data.
        """
        level = float(level)
        new_affine = self.affine * affine.Affine.scale(1 / level)
        if boundary_treatment == "round":
            new_n_cols = int(self.n_cols * level)
            new_n_rows = int(self.n_rows * level)
        elif boundary_treatment == "shrink":
            new_n_cols = int(floor(self.n_cols * level))
            new_n_rows = int(floor(self.n_rows * level))
        elif boundary_treatment == "expand":
            new_n_cols = int(ceil(self.n_cols * level))
            new_n_rows = int(ceil(self.n_rows * level))
        else:
            raise ValueError(f"Unknown `boundary_treatment` {boundary_treatment}")

        new_meta = self.update(
            geotrans=affine.Affine.to_gdal(new_affine),
            n_cols=new_n_cols,
            n_rows=new_n_rows,
        )
        return new_meta


class SourceWindow(NamedTuple):
    """Source Window in image coordinate space.

    Parameters
    ----------
    x_off: int
        Starting x (index of starting column)
    y_off: int
        Starting y (index of starting row)
    x_size: int
        Size of the window in the x direction (number of columns)
    y_size: int
        Size of the window in the y direction (number of rows)
    """

    x_off: int
    y_off: int
    x_size: int
    y_size: int


def calc_bounds_from_geotrans(
    geotrans: Tuple[float, float, float, float, float, float], n_rows: int, n_cols: int
) -> BoundingBox:
    """Calculate bounds from given geotransform tuple and corresponding 2D dimensions.

    Parameters
    ----------
    geotrans : Tuple[float, float, float, float, float, float]
        A 6-element geotransform tuple.
    n_rows : int
        Number of rows.
    n_cols : int
        Number of columns.

    Returns
    -------
    BoundingBox
        Bounding box of the spatial domain.
    """
    x_res, y_res = geotrans[1], geotrans[5]
    x_rot, y_rot = geotrans[2], geotrans[4]
    x_min = geotrans[0]
    x_max = geotrans[0] + n_cols * x_res + n_rows * x_rot
    y_max = geotrans[3]
    y_min = geotrans[3] + n_cols * y_rot + n_rows * y_res
    return BoundingBox(left=x_min, right=x_max, upper=y_max, bottom=y_min)


def calc_window_from_aoi(meta: SceneMeta, aoi: Geom) -> Optional[SourceWindow]:
    """Calculate the window in image pixel space for a given area of interest (AOI) in a geo-raster.

    Parameters
    ----------
    meta : SceneMeta
        Metadata for the geo-raster.
    aoi : Geom
        Area of interest.

    Returns
    -------
    Optional[SourceWindow]
        Calculated window in image pixel space,
        None if such a window does not exist.

    Raises
    ------
    ValueError
        Raised when the coordinate system of the geo-raster is not supported.
    """
    crs = CRS.from_wkt(meta.prj_wkt)
    try:
        ds_epsg = crs.to_epsg()
    except Exception:
        raise ValueError(NOT_COMPATIBLE_SRS_MSG)
    else:
        if ds_epsg not in COMPATIBLE_EPSG:
            raise ValueError(
                "EPSG:{} is not a supported coordinate system.\n{}".format(
                    ds_epsg, NOT_COMPATIBLE_SRS_MSG
                )
            )

    # get the aoi into the crs specified in the meta
    # notice that we require the geometry to be in WGS84 / EPSG:4326,
    # which is a requirement of geojson standard
    # so if already match then skip this
    if crs.to_epsg() != 4326:
        aoi = project_geom(aoi, "epsg:4326", "epsg:{}".format(crs.to_epsg()))

    # formulate a geometry from the image's footprint
    # then calculate intersection between image's footprint and the aoi
    img_bbox = bbox_to_geom(
        calc_bounds_from_geotrans(meta.geotrans, meta.n_rows, meta.n_cols)
    )
    intersection = img_bbox._shapely_geom.intersection(aoi._shapely_geom)
    if intersection.is_empty:
        return None
    min_x, min_y, max_x, max_y = intersection.bounds

    # calculate the the result image's bounds in source images' image space (pixel index)
    fwd = affine.Affine.from_gdal(*(meta.geotrans))
    bwd = ~fwd
    min_i, max_j = bwd * (min_x, min_y)
    max_i, min_j = bwd * (max_x, max_y)

    min_i = int(max(floor(min_i), 0))
    max_i = int(min(ceil(max_i), meta.n_cols - 1))
    min_j = int(max(floor(min_j), 0))
    max_j = int(min(ceil(max_j), meta.n_rows - 1))
    ncols_new = max_i - min_i + 1
    nrows_new = max_j - min_j + 1
    window = SourceWindow(min_i, min_j, ncols_new, nrows_new)
    return window
